package com.stanford.ml.naivebayes;

import java.util.HashMap;
import java.util.HashSet;

/**
 * Data container for training the classifier
 * @author Fatih Sunor
 *
 */
public class TrainingSet {
	
	private HashMap<Integer,Instance> examples;
	private HashSet<ClassLabel> classLabels;
	private HashSet<String> attributes;
	
	/**
	 * Default constructor
	 */
	public TrainingSet() {
	  examples = new HashMap<Integer,Instance>();
	}
	
	/**
	 * Constructor with instances initialized
	 * @param instances
	 */
	public TrainingSet(Instance[] instances) {
		examples = new HashMap<Integer,Instance>();
		classLabels  = new HashSet<ClassLabel>();
		attributes = new HashSet<String>();
		ClassLabel label;
		int exampleCount = 0;
		for (Instance example : instances) {
			examples.put(exampleCount, example);
			label = example.getClassLabel();
			if ( !classLabels.contains(label) ) {
				classLabels.add(label);
			}
			for(Attribute attribute : example.getAttributesWithNominalTag()){
			  if( attribute != null ) {
			    attributes.add(attribute.getName());
			  }
			}
			exampleCount++;
		}
	}

	/**
	 * @return the instanceSet
	 */
	public HashMap<Integer,Instance> getInstances() {
		return examples;
	}

	/**
	 * Returns a particular instance
	 * @param index
	 * @return
	 */
	public Instance getInstance(int index) {
	  return examples.get(index);
	}
	
	/**
	 * @return the size of the instanceSet
	 */
	public int getSize() {
		return examples.size();
	}

	/**
	 * Returns a count of class labels
	 * @return
	 */
	public int getNumberOfClassLabels() {
		return classLabels.size();
	}

	/**
	 * @return the set that contains the class labels
	 */
	public HashSet<ClassLabel> classLabels() {
		return classLabels;
	}
	
	/**
	 * Returns the attributes included
	 * @return the attributes included
	 */
	public HashSet<String> getAttributeNameSet() {
	  return attributes;
	}

}
